# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

from coin.base.datetime_util import to_datetime
from coin.proto import coin_order_pb2
from coin.exchange.base.tick import FixedTickPrice
from coin.strategy.mm.subscription import FeedSubscriptionRequest
from coin.strategy.marketsim.core import spec_pb2, marketsim

from experimental.prophet.graph import graph
from experimental.prophet.graph_runner.fastfeed import FastfeedGraphAdapter

from experimental.prophet.ops.fastfeed import fastfeed_coin
from experimental.prophet.ops import globals, timeseries
from experimental.prophet.ops.timer import timer
from experimental.prophet.ops.debug import printer


def _model(pass_symbol_info,
           ref_symbol_info,
           basis_window,
           machine,
           ref_price_multiplier=1.,
           debug=False):
  ts = globals.global_timestamp()

  with graph.namespace('pass'):
    feed = fastfeed_coin(pass_symbol_info.product, pass_symbol_info.sub_req, machine)
    book = graph.create_node('fastfeed.BookBestLevel', [feed.book])
    pass_ask0_p = book.outputs[1].named_ref('b.ask0_p')
    pass_bid0_p = book.outputs[3].named_ref('b.bid0_p')

  with graph.namespace('ref'):
    feed = fastfeed_coin(ref_symbol_info.product, ref_symbol_info.sub_req, machine)
    book = graph.create_node('fastfeed.BookBestLevel', [feed.book])
    ref_ask0_p = (book.outputs[1] * ref_price_multiplier).named_ref('b.ask0_p')
    ref_bid0_p = (book.outputs[3] * ref_price_multiplier).named_ref('b.bid0_p')

  ask_basis = pass_ask0_p - ref_ask0_p
  bid_basis = pass_bid0_p - ref_bid0_p
  ask_basis_ma = timeseries.time_moving_window(ask_basis,
                                               basis_window).mean().named_ref('ask_basis_ma')
  bid_basis_ma = timeseries.time_moving_window(bid_basis,
                                               basis_window).mean().named_ref('ask_basis_ma')

  proj_ask0_p = (ref_ask0_p + ask_basis_ma).named_ref('proj_ask0_p')
  proj_bid0_p = (ref_bid0_p + bid_basis_ma).named_ref('proj_bid0_p')

  if debug:
    periodic_timer = timer('1m')
    with graph.control_if(ts % (60 * 10**9) == 0):
      printer('%t pass_ask0p=%s ref_ask0_p=%s proj_ask0_p=%s',
              ts,
              pass_ask0_p,
              ref_ask0_p,
              proj_ask0_p)


class Calculator:
  def __init__(self, machine, tick, *args, **kwargs):
    self._tick = FixedTickPrice(tick) if isinstance(tick, float) else tick

    self._graph = graph.Graph()
    with graph.as_default_graph(self._graph):
      _model(*args, machine=machine, **kwargs)

    self._pass_ask0_p = self._graph.get_named_variable('pass/b.ask0_p')
    self._pass_bid0_p = self._graph.get_named_variable('pass/b.bid0_p')
    self._proj_ask0_p = self._graph.get_named_variable('proj_ask0_p')
    self._proj_bid0_p = self._graph.get_named_variable('proj_bid0_p')

    self._adapter = FastfeedGraphAdapter(self._graph, machine)

  @property
  def ready(self):
    return not self._proj_ask0_p.is_null and not self._proj_bid0_p.is_null

  def calculate_price(self, sell_edge, buy_edge, passive=True):
    assert self.ready
    assert -1 < sell_edge < 1
    assert -1 < buy_edge < 1

    proj_ask0_p = self._proj_ask0_p.value
    proj_bid0_p = self._proj_bid0_p.value

    if passive:
      ask0_minus_one = self._tick.get_prev_price(self._pass_ask0_p.value)
      bid0_plus_one = self._tick.get_next_price(self._pass_bid0_p.value)
      sell_price = max(self._tick.ceil(proj_ask0_p * (1. + sell_edge)),
                       ask0_minus_one,
                       bid0_plus_one)
      buy_price = min(self._tick.floor(proj_bid0_p * (1. - buy_edge)),
                      ask0_minus_one,
                      bid0_plus_one)
    else:
      sell_price = self._tick.ceil(proj_ask0_p * (1. + sell_edge))
      buy_price = self._tick.floor(proj_bid0_p * (1. - buy_edge))

    return sell_price, buy_price

  def get_feed_sub_reqs(self, feed_sub_req=None):
    feed_sub_req = feed_sub_req or FeedSubscriptionRequest()
    return self._adapter.get_feed_sub_requests(feed_sub_req)

  @property
  def book_reset_callbacks(self):
    return [self._adapter.on_book_reset]


class PassUnhedgeSimStrategy:
  def __init__(self,
               pass_product_sim_spec,
               pass_symbol_info,
               ref_symbol_info,
               machine,
               edge,
               lot_size,
               min_pos,
               max_pos,
               min_fire_period,
               post_only,
               sim_spec=None,
               name=None,
               **kwargs):
    self._pass_symbol_info = pass_symbol_info
    self._ref_symbol_info = ref_symbol_info

    self._edge = edge
    self._lot_size = lot_size
    self._min_pos = min_pos
    self._max_pos = max_pos
    self._min_fire_period = min_fire_period
    self._post_only = post_only
    self._name = name

    assert str(pass_symbol_info.product) == pass_product_sim_spec.symbol
    self._msim = marketsim.MarketSim(sim_spec=sim_spec, default_machine=machine)
    self._msim.add_product(pass_symbol_info.product, pass_product_sim_spec)
    self._msim.add_execution_callback(self.on_execution)

    self._calc = Calculator(machine=machine,
                            pass_symbol_info=pass_symbol_info,
                            ref_symbol_info=ref_symbol_info,
                            **kwargs)

    self._pos = 0.
    self._last_post_ts = 0

  @property
  def name(self):
    return self._name

  def on_book(self, book):
    if book.timestamp - self._last_post_ts < self._min_fire_period or not self._calc.ready:
      return

    self._msim.cancel_all()

    sell_p, buy_p = self._calc.calculate_price(self._edge, self._edge)
    if self._min_pos < self._pos:
      qty = min(self._pos - self._min_pos, self._lot_size)
      # print('SELL %f %f (%f)' % (sell_p, book.ask0().price, (sell_p-book.ask0().price)/book.ask0().price*10000))
      self._msim.submit(product=self._pass_symbol_info.product,
                        side=coin_order_pb2.SELL_ORDER,
                        price=sell_p,
                        qty=qty,
                        post_only=self._post_only)

    if self._pos < self._max_pos:
      qty = min(self._max_pos - self._pos, self._lot_size)
      # print('BUY %f %f (%f)' % (buy_p, book.bid0().price, (book.bid0().price-buy_p)/book.bid0().price*10000))
      self._msim.submit(product=self._pass_symbol_info.product,
                        side=coin_order_pb2.BUY_ORDER,
                        price=buy_p,
                        qty=qty,
                        post_only=self._post_only)

    self._last_post_ts = book.timestamp

  def on_execution(self, execution):
    self._pos = execution.position

  def get_feed_sub_reqs(self, feed_sub_req=None):
    feed_sub_req = feed_sub_req or FeedSubscriptionRequest()
    self._calc.get_feed_sub_reqs(feed_sub_req)
    self._msim.get_feed_sub_reqs(feed_sub_req)
    return feed_sub_req

  def on_book_reset(self, sub_req, book_builder):
    if sub_req.to_str() == self._pass_symbol_info.sub_req:
      book_builder.subscribe_book(self._pass_symbol_info.product, self.on_book)

  @property
  def book_reset_callbacks(self):
    return self._calc.book_reset_callbacks + [self._msim.on_book_reset, self.on_book_reset]

  @property
  def sim(self):
    return self._msim

  def get_sim_stat(self):
    return self._msim.generate_sim_stat()


'''
import datetime
from absl import app, flags
from coin.feed.fastfeed.native_wrapper import run_from_feed_cache
from experimental.prophet.craps.universe import universe_btc


def main(args):
  machine = 'feed-01.ap-northeast-1.aws'

  symbol_infos = universe_btc(datetime.date(2019, 7, 1))
  fxbtcjpy = symbol_infos[0]
  xbtusd = symbol_infos[1]
  btcusd_quarter = symbol_infos[6]
  btcusd_next_week = symbol_infos[7]

  pass_product = xbtusd
  pass_product_sim_spec = spec_pb2.ProductSpecProto(
      symbol=str(xbtusd.product),
      settlement_currency='BTC',
      inverse=True,
      maker_fee=-0.00025,
      taker_fee=0.00075,
      support_post_only=True,
      fastfeed_machine='feed-01.ap-northeast-1.aws',
      fastfeed_sub_req='Futures.Bitmex')

  ref_product = btcusd_next_week

  strat = PassUnhedgeSimStrategy(
      pass_product_sim_spec,
      pass_product,
      ref_product,
      basis_window='5m',
      machine=machine,
      tick=0.5,
      debug=False)

  run_from_feed_cache(
      strat.get_feed_sub_reqs(),
      strat.book_reset_callbacks,
      datetime.date(2019, 7, 1),
      datetime.date(2019, 7, 2),
      machine=machine)

  print(strat._msim.generate_sim_stat())


if __name__ == '__main__':
  app.run(main)
'''
